clear;
close all;
warning('off','all');
paths;

sample_start = 19781;
sample_end = 20222;

%% load data

% SPF
spf = readtable("spf.csv");
spf.YYYYQ = spf.YEAR * 10 + spf.QUARTER;
spf.spf = [spf.INFPGDP1YR(1:46); spf.INFCPI1YR(47:end)];

% inflation
inflation = readtable("inflation.csv");

% unemployment gap
ue = readtable("unemp.csv"); ue.YYYYQ = year(ue.DATE)*10 + quarter(ue.DATE);
ue.uegap = ue.UNRATE - ue.NROUST;

% oil price
oilprice = readtable("OILPRICE.csv"); oilprice.YYYYQ = year(oilprice.DATE)*10 + quarter(oilprice.DATE);

% MSC

px = readtable("agg_pxeducg.csv");

data = join(px, spf); data = join(data, inflation); 
sample = (data.YYYYQ >= sample_start & data.YYYYQ <= sample_end);
y = [data.spf(sample), data.px1_educg1(sample), data.px1_educg3(sample)];
y5 = [data.INFCPI10YR(sample), data.px5_educg1(sample), data.px5_educg3(sample)];
pi = data.inflation(sample);

save("output/data.mat", 'data', 'y', 'y5', 'pi');

%% Estimation

[T,N] = size(y);

%initials;
xparam = [
    log(0.2);   % v_eta
    log(0.6);   % v_nu1
    log(1);   % v_nu2
    log(0.4);   % v_e       
    ];

hhparam = [
    ones(N,1)*log(0.5);   % v_w
    ones(N,1)*log(1);   % v_e
    ones(N,1)*1;        % theta
    ones(N+1,1)*0.5;        % alpha0
    ones(N+1,1)*0.5     % phi0
    ];

psi0 =[
    6.2; 0.2; -1;    % pi, alpha, phi
    6.7; 5.9; 0.2; -1;
    9.5; 7.2; 0.2; -1;
    8.5; 6.9; 0.2; -1; % s_i, pi_i, alpha_i, phi_i
    5.9; 7.2; 6.9; %LR
    ];

PP0 = diag([1;  % pi
            1;  % alpha
            1;  % phi
            ones(5*N,1)]);% HH

x0 = [xparam; hhparam];

save("output/initial.mat", 'xparam', 'hhparam', 'psi0', 'PP0', 'x0');


options = optimoptions(@fminunc,'Algorithm','quasi-newton','Display','off','MaxFunctionEvaluations',50000,'MaxIterations',5000);
[xg, fg, EXITFLAG,OUTPUT,GRAD,HESSIAN] = fminunc(@(x)obj_main(x, psi0, PP0, y,y5, pi, 0), x0, options);

[~,est] = measurement(xg(1:4), xg(5:end), psi0, PP0, y, y5, pi, 0);


%% Smoothing
S_back(:,1,T) = est.shat(:,1,T);
Sig_back(:,:,T) = est.sig(:,:,T);
std_back(:,1,T) = sqrt(diag(est.sig(:,:,T)));

for t=T-1:-1:1

    [btTp,StTp] = ksmooth(est.H,est.FF(:,:,t+1),est.shat(:,1,t+1),est.sig(:,:,t+1),est.sfor(:,1,t+1),S_back(:,1,t+1),est.omega(:,:,t+1),Sig_back(:,:,t+1));
    S_back(:,1,t) = btTp;
    Sig_back(:,:,t) = StTp;
    std_back(:,1,t) = real(sqrt(diag(StTp)));

end

%% Solution

soln.sd_eta = exp(xg(1)*0.5); 
soln.sd_nu1 = exp(xg(2)*0.5); 
soln.sd_nu2 = exp(xg(3)*0.5); 
soln.sd_e = exp(xg(4)*0.5);

soln.hh_sd_w = exp(xg(5:4+N)*0.5);
soln.hh_sd_e = exp(xg(5+N:4+2*N)*0.5);
soln.hh_theta = xg(5+2*N:4+3*N);

hess = HESSIAN([1,2,3,4,6,7,8,9,10,12,13],[1,2,3,4,6,7,8,9,10,12,13]);
soln.std = sqrt(diag(inv(hess)));

CI.sd_eta = [exp(xg(1)*0.5 - soln.std(1)), exp(xg(1)*0.5 + soln.std(1))];
CI.sd_nu1 = [exp(xg(2)*0.5 - soln.std(2)), exp(xg(2)*0.5 + soln.std(2))];
CI.sd_nu2 = [exp(xg(3)*0.5 - soln.std(3)), exp(xg(3)*0.5 + soln.std(3))];
CI.sd_e = [exp(xg(4)*0.5 - soln.std(4)), exp(xg(4)*0.5 + soln.std(4))];

CI.hh_w_hs = [exp(xg(6)*0.5 - soln.std(5)), exp(xg(6)*0.5 + soln.std(5))];
CI.hh_e_spf = [exp(xg(8)*0.5 - soln.std(7)), exp(xg(8)*0.5 + soln.std(7))];
CI.hh_e_hs = [exp(xg(9)*0.5 - soln.std(8)), exp(xg(9)*0.5 + soln.std(8))];
CI.hh_e_cd = [exp(xg(10)*0.5 - soln.std(9)), exp(xg(10)*0.5 + soln.std(9))];

CI.hh_theta_hs = [soln.hh_theta(2) - 2*soln.std(10), soln.hh_theta(2) + 2*soln.std(10)];
CI.hh_theta_cd = [soln.hh_theta(3) - 2*soln.std(11), soln.hh_theta(3) + 2*soln.std(11)];


alpha = reshape(S_back(2:4:2+4*N,1,:),N+1,T);
phi0 = reshape(S_back(3:4:3+4*N,1,:),N+1,T); phi = 1./(1+0.1*exp(-phi0));
infl_e = reshape(S_back(5:4:5+4*(N-1),1,:),N,T);
infl_5e = reshape(S_back(3+4*N+1:end,1,:),N,T);
signal = reshape(S_back(4:4:4+4*(N-1),1,:),N,T);
pi_actual = reshape(S_back(1,1,:),1,T);

hess = HESSIAN([1,2,3,4,6,7,8,9,10,12,13],[1,2,3,4,6,7,8,9,10,12,13]);
soln.std = sqrt(diag(inv(hess)));



%% Figures

fig;

save("output/soln.mat", 'xg', 'soln', 'CI', 'alpha', 'phi', 'phi0', 'infl_5e', 'infl_e', 'pi_actual', 'signal', 'est');
